from channels.generic.websocket import WebsocketConsumer
from websocket.ssh import SSH
from django.http.request import QueryDict
from assets.models import ServerAssets
from django.utils.six import StringIO
from OPSV2.settings import TMP_DIR
import os
import json
import base64


class WebSSH(WebsocketConsumer):
    message = {'status': 0, 'message': None}
    """
    status:
        0: ssh 连接正常, websocket 正常
        1: 发生未知错误, 关闭 ssh 和 websocket 连接

    message:    
        status 为 1 时, message 为具体的错误信息
        status 为 0 时, message 为 ssh 返回的数据, 前端页面将获取 ssh 返回的数据并写入终端页面
    """

    def connect(self):
        """
        打开 websocket 连接, 通过前端传入的参数尝试连接 ssh 主机
        :return:
        """
        print(self.scope['user'])
        if self.scope["user"].is_anonymous:
            self.close()
        else:
            # await self.channel_layer.group_add(
            #     self.scope['user'].username,
            #     self.channel_name,
            # )
            self.accept()
            query_string = self.scope.get('query_string')
            ssh_args = QueryDict(query_string=query_string, encoding='utf-8')
            width = ssh_args.get('width')
            height = ssh_args.get('height')
            id = ssh_args.get('id')
            s = ServerAssets.objects.get(id=id)
            port = int(float(s.port))
            width = int(width)
            height = int(height)
            host = s.ip
            user = s.username
            passwd = s.password

            self.ssh = SSH(websocker=self, message=self.message)

            ssh_connect_dict = {
                'host': host,
                'user': user,
                'port': port,
                'timeout': 30,
                'pty_width': width,
                'pty_height': height,
                'password': passwd
            }

            self.ssh.connect(**ssh_connect_dict)

    def disconnect(self, close_code):
        try:
            self.ssh.close()
        except:
            pass

    def receive(self, text_data=None, bytes_data=None):
        data = json.loads(text_data)
        if type(data) == dict:
            status = data['status']
            if status == 0:
                data = data['data']
                self.ssh.shell(data)
            else:
                cols = data['cols']
                rows = data['rows']
                self.ssh.resize_pty(cols=cols, rows=rows)
